#include "../camera_models/EquidistantCamera.h"

#include <cmath>
#include <cstdio>
#include <Eigen/Dense>
#include <iomanip>
#include <iostream>
#include <opencv2/calib3d/calib3d.hpp>
#include <opencv2/core/eigen.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include "../gpl/gpl.h"

namespace camodocal {

EquidistantCamera::Parameters::Parameters():
  Camera::Parameters(KANNALA_BRANDT),
  m_k2(0.0),
  m_k3(0.0),
  m_k4(0.0),
  m_k5(0.0),
  m_mu(0.0),
  m_mv(0.0),
  m_u0(0.0),
  m_v0(0.0) {
}

EquidistantCamera::Parameters::Parameters(const std::string& cameraName,
                                          int                w,
                                          int                h,
                                          double             k2,
                                          double             k3,
                                          double             k4,
                                          double             k5,
                                          double             mu,
                                          double             mv,
                                          double             u0,
                                          double             v0):
  Camera::Parameters(KANNALA_BRANDT, cameraName, w, h),
  m_k2(k2),
  m_k3(k3),
  m_k4(k4),
  m_k5(k5),
  m_mu(mu),
  m_mv(mv),
  m_u0(u0),
  m_v0(v0) {
}

double&
EquidistantCamera::Parameters::k2(void) {
  return m_k2;
}

double&
EquidistantCamera::Parameters::k3(void) {
  return m_k3;
}

double&
EquidistantCamera::Parameters::k4(void) {
  return m_k4;
}

double&
EquidistantCamera::Parameters::k5(void) {
  return m_k5;
}

double&
EquidistantCamera::Parameters::mu(void) {
  return m_mu;
}

double&
EquidistantCamera::Parameters::mv(void) {
  return m_mv;
}

double&
EquidistantCamera::Parameters::u0(void) {
  return m_u0;
}

double&
EquidistantCamera::Parameters::v0(void) {
  return m_v0;
}

double
EquidistantCamera::Parameters::k2(void) const {
  return m_k2;
}

double
EquidistantCamera::Parameters::k3(void) const {
  return m_k3;
}

double
EquidistantCamera::Parameters::k4(void) const {
  return m_k4;
}

double
EquidistantCamera::Parameters::k5(void) const {
  return m_k5;
}

double
EquidistantCamera::Parameters::mu(void) const {
  return m_mu;
}

double
EquidistantCamera::Parameters::mv(void) const {
  return m_mv;
}

double
EquidistantCamera::Parameters::u0(void) const {
  return m_u0;
}

double
EquidistantCamera::Parameters::v0(void) const {
  return m_v0;
}

bool EquidistantCamera::Parameters::readFromYamlFile(const std::string& filename) {
  cv::FileStorage fs(filename, cv::FileStorage::READ);

  if(!fs.isOpened()) {
    return false;
  }

  if(!fs["model_type"].isNone()) {
    std::string sModelType;
    fs["model_type"] >> sModelType;

    if(sModelType.compare("KANNALA_BRANDT") != 0) {
      return false;
    }
  }

  m_modelType = KANNALA_BRANDT;
  fs["camera_name"] >> m_cameraName;
  m_imageWidth  = static_cast<int>(fs["image_width"]);
  m_imageHeight = static_cast<int>(fs["image_height"]);

  cv::FileNode n = fs["projection_parameters"];
  m_k2           = static_cast<double>(n["k2"]);
  m_k3           = static_cast<double>(n["k3"]);
  m_k4           = static_cast<double>(n["k4"]);
  m_k5           = static_cast<double>(n["k5"]);
  m_mu           = static_cast<double>(n["mu"]);
  m_mv           = static_cast<double>(n["mv"]);
  m_u0           = static_cast<double>(n["u0"]);
  m_v0           = static_cast<double>(n["v0"]);

  return true;
}

void EquidistantCamera::Parameters::writeToYamlFile(const std::string& filename) const {
  cv::FileStorage fs(filename, cv::FileStorage::WRITE);

  fs << "model_type"
     << "KANNALA_BRANDT";
  fs << "camera_name" << m_cameraName;
  fs << "image_width" << m_imageWidth;
  fs << "image_height" << m_imageHeight;

  // projection: k2, k3, k4, k5, mu, mv, u0, v0
  fs << "projection_parameters";
  fs << "{"
     << "k2" << m_k2
     << "k3" << m_k3
     << "k4" << m_k4
     << "k5" << m_k5
     << "mu" << m_mu
     << "mv" << m_mv
     << "u0" << m_u0
     << "v0" << m_v0 << "}";

  fs.release();
}

EquidistantCamera::Parameters&
EquidistantCamera::Parameters::operator=(const EquidistantCamera::Parameters& other) {
  if(this != &other) {
    m_modelType   = other.m_modelType;
    m_cameraName  = other.m_cameraName;
    m_imageWidth  = other.m_imageWidth;
    m_imageHeight = other.m_imageHeight;
    m_k2          = other.m_k2;
    m_k3          = other.m_k3;
    m_k4          = other.m_k4;
    m_k5          = other.m_k5;
    m_mu          = other.m_mu;
    m_mv          = other.m_mv;
    m_u0          = other.m_u0;
    m_v0          = other.m_v0;
  }

  return *this;
}

std::ostream&
operator<<(std::ostream& out, const EquidistantCamera::Parameters& params) {
  out << "Camera Parameters:" << std::endl;
  out << "    model_type "
      << "KANNALA_BRANDT" << std::endl;
  out << "   camera_name " << params.m_cameraName << std::endl;
  out << "   image_width " << params.m_imageWidth << std::endl;
  out << "  image_height " << params.m_imageHeight << std::endl;

  // projection: k2, k3, k4, k5, mu, mv, u0, v0
  out << "Projection Parameters" << std::endl;
  out << "            k2 " << params.m_k2 << std::endl
      << "            k3 " << params.m_k3 << std::endl
      << "            k4 " << params.m_k4 << std::endl
      << "            k5 " << params.m_k5 << std::endl
      << "            mu " << params.m_mu << std::endl
      << "            mv " << params.m_mv << std::endl
      << "            u0 " << params.m_u0 << std::endl
      << "            v0 " << params.m_v0 << std::endl;

  return out;
}

EquidistantCamera::EquidistantCamera():
  m_inv_K11(1.0),
  m_inv_K13(0.0),
  m_inv_K22(1.0),
  m_inv_K23(0.0) {
}

EquidistantCamera::EquidistantCamera(const std::string& cameraName,
                                     int                imageWidth,
                                     int                imageHeight,
                                     double             k2,
                                     double             k3,
                                     double             k4,
                                     double             k5,
                                     double             mu,
                                     double             mv,
                                     double             u0,
                                     double             v0):
  mParameters(cameraName, imageWidth, imageHeight, k2, k3, k4, k5, mu, mv, u0, v0) {
  // Inverse camera projection matrix parameters
  m_inv_K11 = 1.0 / mParameters.mu();
  m_inv_K13 = -mParameters.u0() / mParameters.mu();
  m_inv_K22 = 1.0 / mParameters.mv();
  m_inv_K23 = -mParameters.v0() / mParameters.mv();
}

EquidistantCamera::EquidistantCamera(const EquidistantCamera::Parameters& params):
  mParameters(params) {
  // Inverse camera projection matrix parameters
  m_inv_K11 = 1.0 / mParameters.mu();
  m_inv_K13 = -mParameters.u0() / mParameters.mu();
  m_inv_K22 = 1.0 / mParameters.mv();
  m_inv_K23 = -mParameters.v0() / mParameters.mv();
}

Camera::ModelType
EquidistantCamera::modelType(void) const {
  return mParameters.modelType();
}

const std::string&
EquidistantCamera::cameraName(void) const {
  return mParameters.cameraName();
}

int EquidistantCamera::imageWidth(void) const {
  return mParameters.imageWidth();
}

int EquidistantCamera::imageHeight(void) const {
  return mParameters.imageHeight();
}

void EquidistantCamera::estimateIntrinsics(const cv::Size&                              boardSize,
                                           const std::vector<std::vector<cv::Point3f>>& objectPoints,
                                           const std::vector<std::vector<cv::Point2f>>& imagePoints) {
  Parameters params = getParameters();

  double u0 = params.imageWidth() / 2.0;
  double v0 = params.imageHeight() / 2.0;

  double minReprojErr = std::numeric_limits<double>::max();

  std::vector<cv::Mat> rvecs, tvecs;
  rvecs.assign(objectPoints.size(), cv::Mat());
  tvecs.assign(objectPoints.size(), cv::Mat());

  params.k2() = 0.0;
  params.k3() = 0.0;
  params.k4() = 0.0;
  params.k5() = 0.0;
  params.u0() = u0;
  params.v0() = v0;

  // Initialize focal length
  // C. Hughes, P. Denny, M. Glavin, and E. Jones,
  // Equidistant Fish-Eye Calibration and Rectification by Vanishing Point
  // Extraction, PAMI 2010
  // Find circles from rows of chessboard corners, and for each pair
  // of circles, find vanishing points: v1 and v2.
  // f = ||v1 - v2|| / PI;
  double f0 = 0.0;
  for(size_t i = 0; i < imagePoints.size(); ++i) {
    std::vector<Eigen::Vector2d> center(boardSize.height);
    std::vector<double>          radius(boardSize.height);
    for(int r = 0; r < boardSize.height; ++r) {
      std::vector<cv::Point2d> circle;
      for(int c = 0; c < boardSize.width; ++c) {
        circle.push_back(imagePoints.at(i).at(r * boardSize.width + c));
      }

      fitCircle(circle, center[r](0), center[r](1), radius[r]);
    }

    for(int j = 0; j < boardSize.height; ++j) {
      for(int k = j + 1; k < boardSize.height; ++k) {
        // find distance between pair of vanishing points which
        // correspond to intersection points of 2 circles
        std::vector<cv::Point2d> ipts;
        ipts = intersectCircles(center[j](0), center[j](1), radius[j], center[k](0), center[k](1), radius[k]);

        if(ipts.size() < 2) {
          continue;
        }

        double f = cv::norm(ipts.at(0) - ipts.at(1)) / M_PI;

        params.mu() = f;
        params.mv() = f;

        setParameters(params);

        for(size_t l = 0; l < objectPoints.size(); ++l) {
          estimateExtrinsics(objectPoints.at(l), imagePoints.at(l), rvecs.at(l), tvecs.at(l));
        }

        double reprojErr = reprojectionError(objectPoints, imagePoints, rvecs, tvecs, cv::noArray());

        if(reprojErr < minReprojErr) {
          minReprojErr = reprojErr;
          f0           = f;
        }
      }
    }
  }

  if(f0 <= 0.0 && minReprojErr >= std::numeric_limits<double>::max()) {
    std::cout << "[" << params.cameraName() << "] "
              << "# INFO: kannala-Brandt model fails with given data. " << std::endl;

    return;
  }

  params.mu() = f0;
  params.mv() = f0;

  setParameters(params);
}

/**
 * \brief Lifts a point from the image plane to the unit sphere
 *
 * \param p image coordinates
 * \param P coordinates of the point on the sphere
 */
void EquidistantCamera::liftSphere(const Eigen::Vector2d& p, Eigen::Vector3d& P) const {
  liftProjective(p, P);
}

/** 
 * \brief Lifts a point from the image plane to its projective ray
 *
 * \param p image coordinates
 * \param P coordinates of the projective ray
 */
void EquidistantCamera::liftProjective(const Eigen::Vector2d& p, Eigen::Vector3d& P) const {
  // Lift points to normalised plane
  Eigen::Vector2d p_u;
  p_u << m_inv_K11 * p(0) + m_inv_K13,
    m_inv_K22 * p(1) + m_inv_K23;

  // Obtain a projective ray
  double theta, phi;
  backprojectSymmetric(p_u, theta, phi);

  P(0) = sin(theta) * cos(phi);
  P(1) = sin(theta) * sin(phi);
  P(2) = cos(theta);
}

/** 
 * \brief Project a 3D point (\a x,\a y,\a z) to the image plane in (\a u,\a v)
 *
 * \param P 3D point coordinates
 * \param p return value, contains the image point coordinates
 */
void EquidistantCamera::spaceToPlane(const Eigen::Vector3d& P, Eigen::Vector2d& p) const {
  double theta = acos(P(2) / P.norm());
  double phi   = atan2(P(1), P(0));

  Eigen::Vector2d p_u = r(mParameters.k2(), mParameters.k3(), mParameters.k4(), mParameters.k5(), theta) * Eigen::Vector2d(cos(phi), sin(phi));

  // Apply generalised projection matrix
  p << mParameters.mu() * p_u(0) + mParameters.u0(),
    mParameters.mv() * p_u(1) + mParameters.v0();
}

/**
 * \brief Project a 3D point to the image plane and calculate Jacobian
 *
 * \param P 3D point coordinates
 * \param p return value, contains the image point coordinates
 */
void EquidistantCamera::spaceToPlane(const Eigen::Vector3d& P, Eigen::Vector2d& p, Eigen::Matrix<double, 2, 3>& J) const {
  double theta = acos(P(2) / P.norm());
  double phi   = atan2(P(1), P(0));

  Eigen::Vector2d p_u = r(mParameters.k2(), mParameters.k3(), mParameters.k4(), mParameters.k5(), theta) * Eigen::Vector2d(cos(phi), sin(phi));

  // Apply generalised projection matrix
  p << mParameters.mu() * p_u(0) + mParameters.u0(),
    mParameters.mv() * p_u(1) + mParameters.v0();
}

/** 
 * \brief Projects an undistorted 2D point p_u to the image plane
 *
 * \param p_u 2D point coordinates
 * \return image point coordinates
 */
void EquidistantCamera::undistToPlane(const Eigen::Vector2d& p_u, Eigen::Vector2d& p) const {
  //    Eigen::Vector2d p_d;
  //
  //    if (m_noDistortion)
  //    {
  //        p_d = p_u;
  //    }
  //    else
  //    {
  //        // Apply distortion
  //        Eigen::Vector2d d_u;
  //        distortion(p_u, d_u);
  //        p_d = p_u + d_u;
  //    }
  //
  //    // Apply generalised projection matrix
  //    p << mParameters.gamma1() * p_d(0) + mParameters.u0(),
  //         mParameters.gamma2() * p_d(1) + mParameters.v0();
}

void EquidistantCamera::initUndistortMap(cv::Mat& map1, cv::Mat& map2, double fScale) const {
  cv::Size imageSize(mParameters.imageWidth(), mParameters.imageHeight());

  cv::Mat mapX = cv::Mat::zeros(imageSize, CV_32F);
  cv::Mat mapY = cv::Mat::zeros(imageSize, CV_32F);

  for(int v = 0; v < imageSize.height; ++v) {
    for(int u = 0; u < imageSize.width; ++u) {
      double mx_u = m_inv_K11 / fScale * u + m_inv_K13 / fScale;
      double my_u = m_inv_K22 / fScale * v + m_inv_K23 / fScale;

      double theta, phi;
      backprojectSymmetric(Eigen::Vector2d(mx_u, my_u), theta, phi);

      Eigen::Vector3d P;
      P << sin(theta) * cos(phi), sin(theta) * sin(phi), cos(theta);

      Eigen::Vector2d p;
      spaceToPlane(P, p);

      mapX.at<float>(v, u) = p(0);
      mapY.at<float>(v, u) = p(1);
    }
  }

  cv::convertMaps(mapX, mapY, map1, map2, CV_32FC1, false);
}

cv::Mat
EquidistantCamera::initUndistortRectifyMap(cv::Mat& map1, cv::Mat& map2, float fx, float fy, cv::Size imageSize, float cx, float cy, cv::Mat rmat) const {
  if(imageSize == cv::Size(0, 0)) {
    imageSize = cv::Size(mParameters.imageWidth(), mParameters.imageHeight());
  }

  cv::Mat mapX = cv::Mat::zeros(imageSize.height, imageSize.width, CV_32F);
  cv::Mat mapY = cv::Mat::zeros(imageSize.height, imageSize.width, CV_32F);

  Eigen::Matrix3f K_rect;

  if(cx == -1.0f && cy == -1.0f) {
    K_rect << fx, 0, imageSize.width / 2,
      0, fy, imageSize.height / 2,
      0, 0, 1;
  }
  else {
    K_rect << fx, 0, cx,
      0, fy, cy,
      0, 0, 1;
  }

  if(fx == -1.0f || fy == -1.0f) {
    K_rect(0, 0) = mParameters.mu();
    K_rect(1, 1) = mParameters.mv();
  }

  Eigen::Matrix3f K_rect_inv = K_rect.inverse();

  Eigen::Matrix3f R, R_inv;
  cv::cv2eigen(rmat, R);
  R_inv = R.inverse();

  for(int v = 0; v < imageSize.height; ++v) {
    for(int u = 0; u < imageSize.width; ++u) {
      Eigen::Vector3f xo;
      xo << u, v, 1;

      Eigen::Vector3f uo = R_inv * K_rect_inv * xo;

      Eigen::Vector2d p;
      spaceToPlane(uo.cast<double>(), p);

      mapX.at<float>(v, u) = p(0);
      mapY.at<float>(v, u) = p(1);
    }
  }

  cv::convertMaps(mapX, mapY, map1, map2, CV_32FC1, false);

  cv::Mat K_rect_cv;
  cv::eigen2cv(K_rect, K_rect_cv);
  return K_rect_cv;
}

int EquidistantCamera::parameterCount(void) const {
  return 8;
}

const EquidistantCamera::Parameters&
EquidistantCamera::getParameters(void) const {
  return mParameters;
}

void EquidistantCamera::setParameters(const EquidistantCamera::Parameters& parameters) {
  mParameters = parameters;

  // Inverse camera projection matrix parameters
  m_inv_K11 = 1.0 / mParameters.mu();
  m_inv_K13 = -mParameters.u0() / mParameters.mu();
  m_inv_K22 = 1.0 / mParameters.mv();
  m_inv_K23 = -mParameters.v0() / mParameters.mv();
}

void EquidistantCamera::readParameters(const std::vector<double>& parameterVec) {
  if(parameterVec.size() != parameterCount()) {
    return;
  }

  Parameters params = getParameters();

  params.k2() = parameterVec.at(0);
  params.k3() = parameterVec.at(1);
  params.k4() = parameterVec.at(2);
  params.k5() = parameterVec.at(3);
  params.mu() = parameterVec.at(4);
  params.mv() = parameterVec.at(5);
  params.u0() = parameterVec.at(6);
  params.v0() = parameterVec.at(7);

  setParameters(params);
}

void EquidistantCamera::writeParameters(std::vector<double>& parameterVec) const {
  parameterVec.resize(parameterCount());
  parameterVec.at(0) = mParameters.k2();
  parameterVec.at(1) = mParameters.k3();
  parameterVec.at(2) = mParameters.k4();
  parameterVec.at(3) = mParameters.k5();
  parameterVec.at(4) = mParameters.mu();
  parameterVec.at(5) = mParameters.mv();
  parameterVec.at(6) = mParameters.u0();
  parameterVec.at(7) = mParameters.v0();
}

void EquidistantCamera::writeParametersToYamlFile(const std::string& filename) const {
  mParameters.writeToYamlFile(filename);
}

std::string
EquidistantCamera::parametersToString(void) const {
  std::ostringstream oss;
  oss << mParameters;

  return oss.str();
}

void EquidistantCamera::fitOddPoly(const std::vector<double>& x, const std::vector<double>& y, int n, std::vector<double>& coeffs) const {
  std::vector<int> pows;
  for(int i = 1; i <= n; i += 2) {
    pows.push_back(i);
  }

  Eigen::MatrixXd X(x.size(), pows.size());
  Eigen::MatrixXd Y(y.size(), 1);
  for(size_t i = 0; i < x.size(); ++i) {
    for(size_t j = 0; j < pows.size(); ++j) {
      X(i, j) = pow(x.at(i), pows.at(j));
    }
    Y(i, 0) = y.at(i);
  }

  Eigen::MatrixXd A = (X.transpose() * X).inverse() * X.transpose() * Y;

  coeffs.resize(A.rows());
  for(int i = 0; i < A.rows(); ++i) {
    coeffs.at(i) = A(i, 0);
  }
}

void EquidistantCamera::backprojectSymmetric(const Eigen::Vector2d& p_u,
                                             double&                theta,
                                             double&                phi) const {
  double tol      = 1e-10;
  double p_u_norm = p_u.norm();

  if(p_u_norm < 1e-10) {
    phi = 0.0;
  }
  else {
    phi = atan2(p_u(1), p_u(0));
  }

  int npow = 9;
  if(mParameters.k5() == 0.0) {
    npow -= 2;
  }
  if(mParameters.k4() == 0.0) {
    npow -= 2;
  }
  if(mParameters.k3() == 0.0) {
    npow -= 2;
  }
  if(mParameters.k2() == 0.0) {
    npow -= 2;
  }

  Eigen::MatrixXd coeffs(npow + 1, 1);
  coeffs.setZero();
  coeffs(0) = -p_u_norm;
  coeffs(1) = 1.0;

  if(npow >= 3) {
    coeffs(3) = mParameters.k2();
  }
  if(npow >= 5) {
    coeffs(5) = mParameters.k3();
  }
  if(npow >= 7) {
    coeffs(7) = mParameters.k4();
  }
  if(npow >= 9) {
    coeffs(9) = mParameters.k5();
  }

  if(npow == 1) {
    theta = p_u_norm;
  }
  else {
    // Get eigenvalues of companion matrix corresponding to polynomial.
    // Eigenvalues correspond to roots of polynomial.
    Eigen::MatrixXd A(npow, npow);
    A.setZero();
    A.block(1, 0, npow - 1, npow - 1).setIdentity();
    A.col(npow - 1) = -coeffs.block(0, 0, npow, 1) / coeffs(npow);

    Eigen::EigenSolver<Eigen::MatrixXd> es(A);
    Eigen::MatrixXcd                    eigval = es.eigenvalues();

    std::vector<double> thetas;
    for(int i = 0; i < eigval.rows(); ++i) {
      if(fabs(eigval(i).imag()) > tol) {
        continue;
      }

      double t = eigval(i).real();

      if(t < -tol) {
        continue;
      }
      else if(t < 0.0) {
        t = 0.0;
      }

      thetas.push_back(t);
    }

    if(thetas.empty()) {
      theta = p_u_norm;
    }
    else {
      theta = *std::min_element(thetas.begin(), thetas.end());
    }
  }
}

} // namespace camodocal
